
#include "simulator.h"

/*
 * File:   mss.c
 * Author: Samik Saha (samiksaha88@gmail.com)
 * Description: This program implements a Mobile Support Station. The
 *              program opens a tcp server in one process for MHs and MSSs
 *              to connect to and in another process it tries to connect to
 *              another MSS.
 * Compiler: gcc
 *
 * Created on 23 November, 2008, 8:58 PM
 */
#include "mss.h"

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <string.h>
#include <pthread.h>
#include <sys/shm.h>
#include <fcntl.h>
#include <sys/ipc.h>
#include <sys/msg.h>

/* Global Variables */

Mobile_support_station mss; //this MSS

int msqid; //Message Queue ID

int *hs_of; //Shared Memory pointer (routing table)

/* Data Structures required for Checkpointing */
int mess_record[NO_MHS];
int ack_record[NO_MHS][NO_MESSAGES];
Message *snt_mess_buff[NO_MESSAGES];
Message *receivemess_buff[NO_MESSAGES];
Message *old_snt_mess_buff[NO_MESSAGES];
Message *old_receive_mess_buf[NO_MESSAGES];
int csn[NO_MHS];

/*
 * Main function
 */
int main(int argc, char** argv) {
  int server_mss;
  int i;

  /* Variables for Message Queue */
  int msgflg = IPC_CREAT | 0666;
  key_t mq_key;

  /* variables for Shared memory */
  int shmid;
  key_t key;

  /* Check for correct program arguments */
  if (argc < 2) {
    printf("Usage:%s MSS-ID\n\n", argv[0]);
    exit(1);
  }

  /* Convert and assign current MSS's id from command line args */
  mss.id = atoi(argv[1]);

  /*~~~~~~~~~~ Create a Shared Memory ~~~~~~~~~~~*/
  key = 5678 + mss.id; //memory segment name
  shmid = shmget(key, 50*sizeof(int), IPC_CREAT | 0666);
  if (shmid < 0) {
    perror("shmget");
    exit(EXIT_FAILURE);
  }
  /* Now we attach the segment to our data space.*/
  if ((hs_of = shmat(shmid, NULL, 0)) == (int *) - 1) {
    perror("shmat");
    exit(EXIT_FAILURE);
  }
  for (i = 0; i < TOTAL_NO_MHS; i++)
    hs_of[i] = -1;

  /*~~~~~~~~~~~~~~~~~ Create a Message Queue ~~~~~~~~~~~~~~~~~~*/
  mq_key = 1234 + mss.id; //generate a unique key
  if ((msqid = msgget(mq_key, msgflg)) < 0) {
    perror("msgget");
    exit(1);
  }

  pid_t pid = fork(); //create a child proces

  if (pid == 0) {//child process to connect to another MSS
    /* ID OF THE ONLY OTHER MSS */
    server_mss = (mss.id == 0) ? 1 : 0;
    connect_to_mss(server_mss);
    _exit(EXIT_SUCCESS);
  }
  else if (pid > 0) {//parent process
    mss_server();//run MSS Server
    return (EXIT_SUCCESS);
  }
}

int mss_server(){
  /* variables for TCP server */
    int sd, newSd, n, rc;
    struct sockaddr_in servAddr;

    Message msg, mbuf;

    int sender_id, msg_type; //msg properties

    /* build server address structure */
    bzero(&servAddr, sizeof (servAddr));
    servAddr.sin_family = AF_INET;
    servAddr.sin_addr.s_addr = inet_addr("127.0.0.1");
    servAddr.sin_port = htons(49152 + mss.id);

    /* create stream socket */
    sd = socket(AF_INET, SOCK_STREAM, 0);
    if (sd < 0) {
      printf("MSS-%d: CANNOT CTEATE STREAM SOCKET \n", mss.id);
      exit(-1);
    } //else printf("MSS %d: successfully created stream socket \n", mss.id);

    /* bind local port number */
    if (bind(sd, (struct sockaddr *) & servAddr, sizeof (servAddr)) < 0) {
      printf("MSS-%d: CANNOT BIND PORT \n", mss.id);
      exit(-1);
    } //else printf("MSS %d: bound local port successfully\n", mss.id);

    listen(sd, 7); /* specify number of concurrent clients to listen to*/

    while (1) {
      //printf("MSS %d: Waiting for client connection\n", mss.id);

      /* wait for client connection*/
      newSd = accept(sd, 0, 0);

      if (newSd < 0) {
        printf("MSS-%d: CANNOT ACCEPT CONNECTION\n", mss.id);
        close(sd);
        exit(-1);
      } else {
        printf("MSS-%d: Received connection\n", mss.id);
      }

      /* Create a child process to serve each MH or MSS that connected to this
       * MSS.
       */
      if (fork() == 0) {//clild process
        pthread_t thread;
        struct thread_data td;

        td.sd = newSd;
        td.mss_id = mss.id;

        close(sd);

        /* Receive first message from client */
        n = receive_message(newSd, &msg);

        sender_id = ntohl(msg.sender_id);
        msg_type = ntohl(msg.msg_type);

        if (msg_type == MH_CONNECT) {
          /*****************    CONNECTED TO AN MH    **********************/
          printf("MSS-%d: MH-%d connected.\n", mss.id, sender_id);

          /* Update routing information */
          hs_of[sender_id] = mss.id;
          msg.qmtype = TOTAL_NO_MHS + ((mss.id == 0) ? 1 : 0);
          enqueue_message(msqid, &msg);
          sleep(1);

          /* Start receiving thread */
          td.src_node_id = sender_id;
          if (pthread_create(&thread, NULL, receive_msgs_from_mh,
                  (void *) & td) != 0) {
            printf("MSS-%d:UNABLE TO CREATE RECEIVING THREAD.\n", mss.id);
            return EXIT_FAILURE;
          }

          /* Sending loop: Send msgs to the MH when data is available */
          while (1) {
            if (dequeue_message(msqid, sender_id, &msg) == -1) {
              printf("MSS-%d: ERROR! QUITING CONN. WITH MH-%d.\n",
                      mss.id, sender_id);
              break;
            }

            //----------------------------------------------------------------
            msg_type=ntohl(msg.msg_type);
            if (msg_type == MH_DISCONNECT) {
              break;
            }
            //----------------------------------------------------------------

            /* now try to send the msg to the MH*/
            rc = send(newSd, &msg, sizeof(msg), 0);
            if (rc < 0) {
              printf("MSS-%d: CANNOT SEND DATA\n", mss.id);
              close(sd);
              exit(1);
            }
            printf("MSS-%d: msg sent to MH-%d - (%s)\n", mss.id,
                    sender_id, msg.content);
            //sleep(1);
          }//end MH_CONNECT
        } else if (msg_type == MSS_CONNECT) {

          /*******************   CONNECTED TO ANOTHER MSS   ******************/
          printf("MSS-%d: Connected to MSS-%d\n", mss.id, sender_id);

          /* Start thread for receiving from MSS */
          td.src_node_id = sender_id;
          if (pthread_create(&thread, NULL, receive_msgs_from_mss,
                  (void *) & td) != 0) {
            printf("MSS-%d:UNABLE TO CREATE RECEIVING THREAD.\n", mss.id);
            return EXIT_FAILURE;
          }

          /* Start loop for sending to MSS */
          while (1) {
            if (dequeue_message(msqid, TOTAL_NO_MHS + sender_id, &msg) == -1){
              printf("MSS-%d:ERROR! ENDING CONN. WITH MSS-%d",
                      mss.id, sender_id);
              break;
            }
            
            rc = send(newSd, &msg, sizeof(msg), 0);
            if (rc < 0) {
              printf("MSS-%d: CANNOT SEND MSG TO MSS-%d\n", mss.id,
                      sender_id);
              close(sd);
              _exit(EXIT_FAILURE);
            } else
              printf("MSS-%d: msg sent to MSS-%d (%s)\n", mss.id,
                    sender_id, msg.content);
            //sleep(1);
          }//end while loop
        }//end MSS_CONNECT
        /* Wait for receiving thread to end */
        pthread_join(thread, NULL);

        /* close client connection*/
        printf("MSS-%d: MH-%d Disconnected. \n", mss.id, sender_id);
        shutdown(newSd, 2);
        close(newSd);

        _exit(EXIT_SUCCESS);
      }//end child process
    else
      close(newSd);
  }//end while loop
}

int connect_to_mss(int server_mss) {
  pthread_t thread;
  struct thread_data td;

  /* Variables for socket connection */
  int sd, rc;
  struct sockaddr_in servAddr;

  Message msg;

  /* build server address structure */
  bzero((char *) & servAddr, sizeof (servAddr));
  servAddr.sin_family = AF_INET;
  servAddr.sin_addr.s_addr = inet_addr("127.0.0.1");
  servAddr.sin_port = htons(49152 + server_mss);

  /* create stream socket */
  sd = socket(AF_INET, SOCK_STREAM, 0);
  if (sd < 0) {
    printf("MSS-%d:(c) CANNOT CREATE STREAM SOCKET\n", mss.id);
    _exit(EXIT_FAILURE);
  } //else printf("MSS %d: created stream socket\n", mss.id);

  /* connect to server */
  printf("MSS-%d: Trying to connect to MSS-%d...\n", mss.id, server_mss);
  rc = connect(sd, (struct sockaddr *) & servAddr, sizeof (servAddr));
  if (rc < 0) {
    printf("MSS-%d: MSS%d not available\n", mss.id, server_mss);
    _exit(EXIT_FAILURE);
  }

  /* construct CONNECT-TO-MSS msg */
  msg.msg_type = htonl(MSS_CONNECT);
  msg.sender_id = htonl(mss.id);
  strcpy(msg.content, "MSS CONNECT");

  /* send control msg to server */
  rc = send(sd, &msg, sizeof(msg), 0);
  if (rc < 0) {
    printf("MSS-%d: CANNOT SEND CTRL MSG TO MSS-%d\n", mss.id, server_mss);
    close(sd);
    _exit(EXIT_FAILURE);
  }

  printf("MSS-%d: Connected to MSS-%d\n", mss.id, server_mss);

  /* Start receiving thread */
  td.mss_id = mss.id;
  td.sd = sd;
  td.src_node_id = server_mss;
  if (pthread_create(&thread, NULL, receive_msgs_from_mss, (void*) & td) != 0) {
    printf("MSS-%d:UNABLE TO CREATE RECEIVING THREAD.\n", mss.id);
    return EXIT_FAILURE;
  }

  /* Start sending loop */
  while (1) {
    if (dequeue_message(msqid, TOTAL_NO_MHS + server_mss, &msg) == -1) break;

    rc = send(sd, &msg, sizeof(msg), 0);
    if (rc < 0) {
      printf("MSS-%d: CANNOT SEND MSG TO MSS-%d\n", mss.id, server_mss);
      close(sd);
      exit(1);
    } else
      printf("MSS-%d: msg (%s) sent to MSS-%d\n", mss.id, msg.content, server_mss);
    //sleep(1);
  }
  close(sd);
  printf("MSS-%d: Connection closed with MSS-%d.\n", mss.id, server_mss);
}

static void *receive_msgs_from_mh(void *tdarg) {
  struct thread_data *td;
  Message msg;
  int n, msg_type, sender_id, dest_id;
  int other_mss;

  td = (struct thread_data *) tdarg;
  other_mss = (td->mss_id == 0) ? 1 : 0;

  while (1) {
    n = receive_message(td->sd, &msg);
    if (n <= 0) {//Error in receiving
      printf("MSS-%d: ERROR RECEIVING FROM MH-%d\n", td->mss_id, td->src_node_id);
      //hs_of[td->src_node_id] = -1; //Remove entry from Routing Table

      msg.msg_type = htonl(MH_DISCONNECT);
      msg.sender_id = htonl(td->src_node_id);

      msg.qmtype = td->src_node_id;

      enqueue_message(msqid, &msg);

      break;
    }

    //printf("Received msg. size: %d content:%s\n", strlen(msg), msg);
    //msg[n] = '\n';
    msg_type = ntohl(msg.msg_type);
    sender_id = ntohl(msg.sender_id);
    dest_id = ntohl(msg.dest_id);

    switch (msg_type) {
      case MH_DISCONNECT://ctrl msg - MH disconnection
        hs_of[sender_id] = -1; //remove entry for HS
        msg.qmtype = sender_id;
        enqueue_message(msqid, &msg); //--
        return NULL;
        break;
      case COMPT_MSG://for computation msgs
        printf("MSS-%d: Received COMPT msg from MH-%d (%s)\n", td->mss_id,
                sender_id, msg.content);
        //printf("hs of %d = %d\n", dest_id, hs_of[dest_id]);
        /* if destination MH is under current MSS or unknown */
        if (hs_of[dest_id] == -1 || hs_of[dest_id] == (td->mss_id)) {
          /* then queue msg for that MH (send to pipe) */
          msg.qmtype = dest_id;
          enqueue_message(msqid, &msg);
        } else { /* else queue them for the other MSS */
          msg.qmtype = TOTAL_NO_MHS + other_mss;
          enqueue_message(msqid, &msg);
        }
        break;
      case ACK_MSG:
        printf("MSS-%d: Received ACK msg from MH-%d\n", td->mss_id,
                sender_id);
        if (hs_of[dest_id] == -1 || hs_of[dest_id] == (td->mss_id)) {
          /* then queue msg for that MH (send to pipe) */
          msg.qmtype = dest_id;
          enqueue_message(msqid, &msg);
        } else { /* else queue them for the other MSS */
          msg.qmtype = TOTAL_NO_MHS + other_mss;
          enqueue_message(msqid, &msg);
        }
        break;
      default:
        printf("MSS-%d: UNKNOWN MESSAGE TYPE FROM MH-%d\n", td->mss_id,
                td->src_node_id);
        break;
    }//end switch
  }//end while
  return NULL;
}//end receive_msgs_from_mh()

static void *receive_msgs_from_mss(void *tdarg) {
  struct thread_data *td;
  Message msg;
  int n, msg_type, sender_id, dest_id;
  int other_mss;

  td = (struct thread_data *) tdarg;
  other_mss = (td->mss_id == 0) ? 1 : 0;

  while (1) {
    n = receive_message(td->sd, &msg);
    if (n <= 0) {//Error in receiving
      printf("MSS-%d:ERROR RCVING FROM MSS-%d\n", td->mss_id, td->src_node_id);
      break;
    }

    msg_type = ntohl(msg.msg_type);
    sender_id = ntohl(msg.sender_id);
    dest_id = ntohl(msg.dest_id);

    printf("from mss: %d %d %d %s\n", msg_type,sender_id, dest_id, msg.content);
    switch (msg_type) {
      case MH_CONNECT://ctrl msg - new MH connection
        printf("MSS-%d: Received MH CONNECT msg from MSS-%d\n", td->mss_id,
                td->src_node_id);

        hs_of[sender_id] = dest_id; //update HS info (routing)
        /*
         * check if there are messeges already queued for this newly
         * connected MH. If so, direct them to the the HS of this MH
         */
        while (dequeue_message_nowait(msqid, sender_id, &msg) != -1) {
          //printf("forwarding msg to MSS-%d (%s)\n", dest_id, msg.content);
          msg.qmtype = TOTAL_NO_MHS + dest_id; //here dest_id is hs id
          enqueue_message(msqid, &msg);
        }

        break;

      case MH_DISCONNECT:
        printf("MSS-%d: Rcvd MH DISCONN.msg from MSS-%d\n", td->mss_id,
                td->src_node_id);
        hs_of[sender_id] = -1; //update HS info (routing)
        break;
      case COMPT_MSG://computation msg
        printf("MSS-%d: Received COMPT msg from MSS-%d (%s)\n", td->mss_id,
                td->src_node_id, msg.content);
        if (hs_of[dest_id] == -1 || hs_of[dest_id] == (td->mss_id)) {
          msg.qmtype = dest_id;
          enqueue_message(msqid, &msg);
        } else { /* else queue them for the other MSS */
          msg.qmtype = TOTAL_NO_MHS + other_mss;
          enqueue_message(msqid, &msg);
        }
        break;
      case ACK_MSG:
        printf("MSS-%d: Received ACK msg from MSS-%d\n", td->mss_id,
                td->src_node_id);
        if (hs_of[dest_id] == -1 || hs_of[dest_id] == (td->mss_id)) {
          msg.qmtype = dest_id;
          enqueue_message(msqid, &msg);
        } else { /* else queue them for the other MSS */
          msg.qmtype = TOTAL_NO_MHS + other_mss;
          enqueue_message(msqid, &msg);
        }
        break;
      default:
        printf("MSS-%d: UNKNOWN MESSAGE TYPE FROM MSS-%d\n",
                td->mss_id, td->src_node_id);
        break;
    }//end switch
  }//end while
}//end receive_msgs_from_mss

